"""
根据LabelMe标注文件生成黑白掩膜图像
"""
import os
import json
from PIL import Image, ImageDraw
from tqdm import tqdm  # 已经导入


def generate_masks(image_dir, json_dir, output_dir):
    """
    根据LabelMe标注文件生成黑白掩膜图像

    参数:
    image_dir: 原始图片文件夹路径
    json_dir: JSON标注文件文件夹路径
    output_dir: 掩膜输出文件夹路径
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    # 获取所有图片文件列表
    img_files = [f for f in os.listdir(image_dir)
                 if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp'))]

    # 使用tqdm包装循环
    for img_name in tqdm(img_files, desc="生成掩膜"):
        img_path = os.path.join(image_dir, img_name)
        base_name = os.path.splitext(img_name)[0]
        json_path = os.path.join(json_dir, f"{base_name}.json")

        # 检查JSON文件是否存在
        if not os.path.exists(json_path):
            print(f"\n警告：找不到 {img_name} 对应的JSON文件，已跳过")  # 加\n避免干扰进度条
            continue

        try:
            # 打开原始图片获取尺寸
            with Image.open(img_path) as img:
                width, height = img.size

            # 创建全黑画布（单通道）
            mask = Image.new('L', (width, height), 0)
            draw = ImageDraw.Draw(mask)

            # 读取JSON标注数据
            with open(json_path, 'r', encoding='utf-8') as f:
                data = json.load(f)

            # 绘制所有标注形状
            for shape in data.get('shapes', []):
                shape_type = shape.get('shape_type', 'polygon')
                points = shape.get('points', [])

                if shape_type == 'polygon':
                    if len(points) < 3:
                        continue
                    # 将浮点坐标转换为整数
                    polygon = [(int(x), int(y)) for x, y in points]
                    draw.polygon(polygon, fill=255)

                elif shape_type == 'rectangle':
                    if len(points) != 2:
                        continue
                    # 提取并排序坐标点
                    x_coords = sorted([points[0][0], points[1][0]])
                    y_coords = sorted([points[0][1], points[1][1]])
                    draw.rectangle(
                        [x_coords[0], y_coords[0], x_coords[1], y_coords[1]],
                        fill=255
                    )

            # 保存掩膜文件
            output_path = os.path.join(output_dir, f"{base_name}.png")
            mask.save(output_path)

        except Exception as e:
            print(f"\n处理 {img_name} 时发生错误：{str(e)}")  # 加\n避免干扰进度条


if __name__ == "__main__":
    # 原图(大图)路径
    IMAGE_DIR = "train_img_big"
    # JSON标注文件路径
    JSON_DIR = "json"
    # 输出路径(输出为黑白掩码图[大图])
    OUTPUT_DIR = "train_lab_big"

    generate_masks(IMAGE_DIR, JSON_DIR, OUTPUT_DIR)